import requests
from config import OLLAMA_BASE_URL, EMBEDDING_MODEL
from typing import List, Union
from services.ollama_service import OllamaService
import numpy as np

class EmbeddingService:
    """统一的嵌入服务，支持不同的嵌入模型"""
    
    def __init__(self):
        self.ollama_service = OllamaService()
        
    def get_embeddings(self, texts: Union[str, List[str]], model_name: str) -> Union[List[float], List[List[float]]]:
        """获取文本的向量表示
        
        Args:
            texts: 单个文本或文本列表
            model_name: 要使用的嵌入模型名称
            
        Returns:
            如果输入是单个文本，返回一个向量
            如果输入是文本列表，返回向量列表
        """
        # 确保texts是列表
        if isinstance(texts, str):
            texts = [texts]
            single_text = True
        else:
            single_text = False
            
        try:
            # 使用Ollama获取嵌入向量
            embeddings = [
                self.ollama_service.get_embeddings(text, model_name) 
                for text in texts
            ]
            
            # 如果是单个文本输入，返回单个向量
            return embeddings[0] if single_text else embeddings
            
        except Exception as e:
            raise Exception(f"获取嵌入向量失败: {str(e)}")